#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PX4飞控控制模拟程序
功能：
1. 模拟PX4飞控系统的基本功能
2. 从后端获取任务和航线信息
3. 执行航线飞行控制
4. 实时反馈无人机状态
5. 支持基本的飞控指令
"""

import asyncio
import time
import json
import random
import requests
import math
from datetime import datetime
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("px4_simulator.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("PX4Simulator")

class PX4ControlSimulator:
    def __init__(self):
        # 后端API地址
        self.base_url = "http://localhost:8000"
        # 模拟的无人机ID
        self.drone_id = 1
        # 当前执行的任务
        self.current_task = None
        # 当前航线信息
        self.current_route = None
        # 当前航点索引
        self.current_waypoint_index = 0
        
        # PX4飞控状态
        self.px4_status = {
            "latitude": 0.0,
            "longitude": 0.0,
            "altitude": 0.0,
            "relative_alt": 0.0,
            "speed_north": 0.0,
            "speed_east": 0.0,
            "speed_down": 0.0,
            "ground_speed": 0.0,
            "heading": 0.0,
            "roll": 0.0,
            "pitch": 0.0,
            "yaw": 0.0,
            "arming_state": "DISARMED",  # ARMED, DISARMED
            "flight_mode": "MANUAL",    # MANUAL, STABILIZED, ALTCTL, POSCTL, AUTO.MISSION, AUTO.LOITER
            "battery_percentage": 100,
            "battery_voltage": 12.6,
            "battery_current": 0.0,
            "gps_satellites": 15,
            "gps_fix_type": 3,
            "armed": False,
            "motors_on": False,
            "mission_item_current": 0,
            "mission_item_total": 0,
            "status_text": "Ready to fly",
            # 内部状态
            "motors_speed": [0, 0, 0, 0],  # 四个电机的速度
            "esc_temperature": [30.0, 30.0, 30.0, 30.0],  # ESC温度
            "cpu_temperature": 45.0,
            "cpu_usage": 30.0,
            "memory_usage": 40.0,
            "vibration_x": 0.1,
            "vibration_y": 0.1,
            "vibration_z": 0.1,
            "airspeed": 0.0,
            "throttle": 0.5
        }
        
        # 模拟参数
        self.update_interval = 0.1  # 100ms更新一次
        self.task_check_interval = 5.0
        self.simulation_speed = 1.0
        self.is_flying = False
        
        # 控制参数
        self.control_params = {
            "position_p": 0.5,  # 位置PID参数
            "position_i": 0.1,
            "position_d": 0.2,
            "altitude_p": 0.6,
            "altitude_i": 0.1,
            "altitude_d": 0.2,
            "heading_p": 0.8,
            "heading_i": 0.05,
            "heading_d": 0.1
        }
        
        # 导航状态
        self.navigation_status = {
            "target_latitude": 0.0,
            "target_longitude": 0.0,
            "target_altitude": 0.0,
            "distance_to_target": 0.0,
            "eta": 0.0,  # 预计到达时间
            "nav_state": "IDLE"  # IDLE, TAKEOFF, WAYPOINT, LOITER, LANDING
        }
    
    def haversine_distance(self, lat1, lon1, lat2, lon2):
        """计算两点之间的距离（米）"""
        R = 6371000  # 地球半径（米）
        phi1 = math.radians(lat1)
        phi2 = math.radians(lat2)
        delta_phi = math.radians(lat2 - lat1)
        delta_lambda = math.radians(lon2 - lon1)
        
        a = math.sin(delta_phi / 2) ** 2 + \
            math.cos(phi1) * math.cos(phi2) * math.sin(delta_lambda / 2) ** 2
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
        
        return R * c
    
    def calculate_bearing(self, lat1, lon1, lat2, lon2):
        """计算两点之间的航向角"""
        dLon = math.radians(lon2 - lon1)
        y = math.sin(dLon) * math.cos(math.radians(lat2))
        x = math.cos(math.radians(lat1)) * math.sin(math.radians(lat2)) - \
            math.sin(math.radians(lat1)) * math.cos(math.radians(lat2)) * math.cos(dLon)
        
        return (math.degrees(math.atan2(y, x)) + 360) % 360
    
    def move_towards(self, lat1, lon1, lat2, lon2, distance):
        """向目标点移动指定距离"""
        total_distance = self.haversine_distance(lat1, lon1, lat2, lon2)
        
        if total_distance == 0:
            return lat1, lon1
        
        # 计算移动比例
        ratio = distance / total_distance
        if ratio > 1:
            ratio = 1
        
        # 计算新位置
        lat = lat1 + (lat2 - lat1) * ratio
        lon = lon1 + (lon2 - lon1) * ratio
        
        return lat, lon
    
    def arm_drone(self):
        """模拟无人机解锁"""
        logger.info("[PX4] 执行解锁操作...")
        self.px4_status["arming_state"] = "ARMING"
        self.px4_status["status_text"] = "Arming motors..."
        time.sleep(0.5)
        
        # 模拟电机启动声音
        for i in range(4):
            self.px4_status["motors_speed"][i] = random.uniform(1000, 1200)
            self.px4_status["esc_temperature"][i] += 1.0
        
        self.px4_status["arming_state"] = "ARMED"
        self.px4_status["armed"] = True
        self.px4_status["motors_on"] = True
        self.px4_status["status_text"] = "Armed and ready"
        logger.info("[PX4] 无人机已解锁")
        return True
    
    def disarm_drone(self):
        """模拟无人机上锁"""
        logger.info("[PX4] 执行上锁操作...")
        self.px4_status["arming_state"] = "DISARMING"
        self.px4_status["status_text"] = "Disarming motors..."
        time.sleep(0.5)
        
        # 模拟电机停止
        for i in range(4):
            self.px4_status["motors_speed"][i] = 0
        
        self.px4_status["arming_state"] = "DISARMED"
        self.px4_status["armed"] = False
        self.px4_status["motors_on"] = False
        self.px4_status["status_text"] = "Disarmed"
        logger.info("[PX4] 无人机已上锁")
        return True
    
    def set_flight_mode(self, mode):
        """设置飞行模式"""
        valid_modes = ["MANUAL", "STABILIZED", "ALTCTL", "POSCTL", "AUTO.MISSION", "AUTO.LOITER"]
        if mode not in valid_modes:
            logger.error(f"[PX4] 不支持的飞行模式: {mode}")
            return False
        
        self.px4_status["flight_mode"] = mode
        self.px4_status["status_text"] = f"Flight mode: {mode}"
        logger.info(f"[PX4] 飞行模式切换至: {mode}")
        return True
    
    def get_pending_tasks(self):
        """从后端获取待执行的任务"""
        try:
            response = requests.get(f"{self.base_url}/api/tasks", params={"status": "pending"})
            if response.status_code == 200:
                data = response.json()
                if data["items"]:
                    # 返回优先级最高的任务
                    return sorted(data["items"], key=lambda x: (x["priority"] == "urgent", x["created_at"]), reverse=True)[0]
            return None
        except Exception as e:
            logger.error(f"获取待执行任务失败: {str(e)}")
            return None
    
    def get_task_routes(self, task_id):
        """获取任务关联的航线"""
        try:
            response = requests.get(f"{self.base_url}/api/tasks/{task_id}")
            if response.status_code == 200:
                task_data = response.json()
                routes = []
                # 按执行顺序排序航线
                for task_route in sorted(task_data.get("task_routes", []), key=lambda x: x.get("order_index", 0)):
                    if task_route.get("route"):
                        routes.append(task_route["route"])
                return routes
            return []
        except Exception as e:
            logger.error(f"获取任务航线失败: {str(e)}")
            return []
    
    def update_task_status(self, task_id, status, progress=0):
        """更新任务状态"""
        try:
            data = {
                "status": status,
                "progress": progress
            }
            # 如果是开始任务，添加开始时间
            if status == "in_progress" and progress == 0:
                data["start_time"] = datetime.now().isoformat()
            # 如果是完成任务，添加结束时间
            elif status == "completed" and progress == 100:
                data["end_time"] = datetime.now().isoformat()
            
            response = requests.put(f"{self.base_url}/api/tasks/{task_id}/status", json=data)
            return response.status_code == 200
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            return False
    
    def send_flight_data(self, task_id):
        """向后台发送飞行数据"""
        try:
            # 构建PX4风格的飞行日志数据
            flight_data = {
                "drone_id": self.drone_id,
                "task_id": task_id,
                "latitude": round(self.px4_status["latitude"], 7),
                "longitude": round(self.px4_status["longitude"], 7),
                "altitude": round(self.px4_status["altitude"], 2),
                "relative_altitude": round(self.px4_status["relative_alt"], 2),
                "ground_speed": round(self.px4_status["ground_speed"], 2),
                "heading": round(self.px4_status["heading"], 2),
                "roll": round(self.px4_status["roll"], 2),
                "pitch": round(self.px4_status["pitch"], 2),
                "yaw": round(self.px4_status["yaw"], 2),
                "battery_percentage": self.px4_status["battery_percentage"],
                "battery_voltage": self.px4_status["battery_voltage"],
                "battery_current": self.px4_status["battery_current"],
                "flight_mode": self.px4_status["flight_mode"],
                "arming_state": self.px4_status["arming_state"],
                "gps_satellites": self.px4_status["gps_satellites"],
                "gps_fix_type": self.px4_status["gps_fix_type"],
                "airspeed": round(self.px4_status["airspeed"], 2),
                "throttle": round(self.px4_status["throttle"], 2),
                "cpu_temperature": round(self.px4_status["cpu_temperature"], 1),
                "vibration_x": round(self.px4_status["vibration_x"], 2),
                "vibration_y": round(self.px4_status["vibration_y"], 2),
                "vibration_z": round(self.px4_status["vibration_z"], 2),
                "mission_item_current": self.px4_status["mission_item_current"],
                "mission_item_total": self.px4_status["mission_item_total"],
                "status_text": self.px4_status["status_text"]
            }
            
            # 发送数据到后端
            response = requests.post(f"{self.base_url}/api/flight_logs", json=flight_data)
            if response.status_code == 200:
                logger.debug(f"[PX4] 飞行数据发送成功: 位置=({flight_data['latitude']}, {flight_data['longitude']}), 高度={flight_data['altitude']}m")
                return True
            else:
                logger.error(f"[PX4] 飞行数据发送失败: {response.status_code} - {response.text}")
                return False
        except Exception as e:
            logger.error(f"[PX4] 发送飞行数据异常: {str(e)}")
            return False
    
    def update_drone_status(self):
        """更新无人机状态"""
        try:
            data = {
                "id": self.drone_id,
                "status": "busy" if self.is_flying else "available",
                "latitude": self.px4_status["latitude"],
                "longitude": self.px4_status["longitude"],
                "altitude": self.px4_status["altitude"],
                "battery_level": self.px4_status["battery_percentage"]
            }
            
            response = requests.put(f"{self.base_url}/api/drones/{self.drone_id}", json=data)
            return response.status_code == 200
        except Exception as e:
            logger.error(f"[PX4] 更新无人机状态失败: {str(e)}")
            return False
    
    def takeoff(self, target_altitude=10.0):
        """模拟起飞过程"""
        logger.info(f"[PX4] 执行起飞指令，目标高度: {target_altitude}m")
        
        # 解锁无人机
        if not self.px4_status["armed"]:
            if not self.arm_drone():
                return False
        
        # 设置为自动模式
        self.set_flight_mode("POSCTL")
        
        # 更新导航状态
        self.navigation_status["nav_state"] = "TAKEOFF"
        self.px4_status["status_text"] = "Taking off..."
        
        # 模拟起飞过程
        start_alt = self.px4_status["altitude"]
        takeoff_time = abs(target_altitude - start_alt) / 2.0  # 假设爬升速度为2m/s
        start_time = time.time()
        
        while time.time() - start_time < takeoff_time:
            # 更新高度
            progress = (time.time() - start_time) / takeoff_time
            new_alt = start_alt + (target_altitude - start_alt) * progress
            self.px4_status["altitude"] = new_alt
            self.px4_status["relative_alt"] = new_alt
            
            # 更新电机速度和温度
            for i in range(4):
                self.px4_status["motors_speed"][i] = 1500 + random.uniform(-50, 50)
                self.px4_status["esc_temperature"][i] = min(70.0, 35.0 + progress * 20)
            
            # 更新油门和电池
            self.px4_status["throttle"] = 0.6 + random.uniform(-0.05, 0.05)
            self.px4_status["battery_percentage"] -= 0.05
            
            # 模拟轻微的姿态变化
            self.px4_status["roll"] = random.uniform(-2, 2)
            self.px4_status["pitch"] = random.uniform(-2, 2)
            
            time.sleep(self.update_interval)
        
        # 稳定在目标高度
        self.px4_status["altitude"] = target_altitude
        self.px4_status["relative_alt"] = target_altitude
        self.px4_status["status_text"] = f"Hovering at {target_altitude}m"
        logger.info(f"[PX4] 起飞完成，当前高度: {target_altitude}m")
        
        return True
    
    def land(self):
        """模拟降落过程"""
        logger.info("[PX4] 执行降落指令")
        
        # 更新导航状态
        self.navigation_status["nav_state"] = "LANDING"
        self.px4_status["status_text"] = "Landing..."
        
        # 设置为降落模式
        self.set_flight_mode("ALTCTL")
        
        # 模拟降落过程
        start_alt = self.px4_status["altitude"]
        landing_time = start_alt / 1.0  # 假设下降速度为1m/s
        start_time = time.time()
        
        while time.time() - start_time < landing_time:
            # 更新高度
            progress = (time.time() - start_time) / landing_time
            new_alt = start_alt * (1 - progress)
            self.px4_status["altitude"] = new_alt
            self.px4_status["relative_alt"] = new_alt
            
            # 更新电机速度和温度
            for i in range(4):
                self.px4_status["motors_speed"][i] = 1300 + random.uniform(-50, 50) * (1 - progress)
                self.px4_status["esc_temperature"][i] = max(35.0, 65.0 - progress * 30)
            
            # 更新油门和电池
            self.px4_status["throttle"] = 0.4 + random.uniform(-0.05, 0.05)
            self.px4_status["battery_percentage"] -= 0.03
            
            # 模拟轻微的姿态变化
            self.px4_status["roll"] = random.uniform(-3, 3)
            self.px4_status["pitch"] = random.uniform(-3, 3)
            
            time.sleep(self.update_interval)
        
        # 到达地面，上锁
        self.px4_status["altitude"] = 0
        self.px4_status["relative_alt"] = 0
        self.disarm_drone()
        self.px4_status["status_text"] = "Landed and disarmed"
        logger.info("[PX4] 降落完成")
        
        return True
    
    def execute_mission_item(self, waypoint):
        """执行单个航点任务"""
        target_lat = waypoint["lat"]
        target_lng = waypoint["lng"]
        target_alt = waypoint["alt"]
        
        logger.info(f"[PX4] 飞向航点 {self.current_waypoint_index + 1}: ({target_lat}, {target_lng}), 高度: {target_alt}m")
        
        # 更新导航状态
        self.navigation_status.update({
            "target_latitude": target_lat,
            "target_longitude": target_lng,
            "target_altitude": target_alt,
            "distance_to_target": self.haversine_distance(
                self.px4_status["latitude"], 
                self.px4_status["longitude"], 
                target_lat, 
                target_lng
            ),
            "nav_state": "WAYPOINT"
        })
        
        self.px4_status["mission_item_current"] = self.current_waypoint_index + 1
        self.px4_status["status_text"] = f"Navigating to waypoint {self.current_waypoint_index + 1}"
        
        # 设置为自动任务模式
        self.set_flight_mode("AUTO.MISSION")
        
        # 飞向目标点
        waypoint_reached = False
        while not waypoint_reached:
            # 计算到目标的距离
            distance = self.haversine_distance(
                self.px4_status["latitude"], 
                self.px4_status["longitude"], 
                target_lat, 
                target_lng
            )
            
            # 更新导航状态
            self.navigation_status["distance_to_target"] = distance
            
            # 如果接近目标航点（5米以内），认为到达
            if distance < 5.0:
                waypoint_reached = True
                break
            
            # 计算航向
            heading = self.calculate_bearing(
                self.px4_status["latitude"], 
                self.px4_status["longitude"], 
                target_lat, 
                target_lng
            )
            
            # 模拟飞行速度（2-8 m/s）
            flight_speed = random.uniform(2.0, 8.0)
            
            # 计算本次移动的距离
            move_distance = flight_speed * self.update_interval * self.simulation_speed
            
            # 更新位置
            new_lat, new_lng = self.move_towards(
                self.px4_status["latitude"], 
                self.px4_status["longitude"], 
                target_lat, 
                target_lng, 
                move_distance
            )
            
            # 平滑更新高度
            altitude_diff = target_alt - self.px4_status["altitude"]
            new_alt = self.px4_status["altitude"] + altitude_diff * 0.05  # 5%的步长接近目标高度
            
            # 更新飞行状态
            self.px4_status.update({
                "latitude": new_lat,
                "longitude": new_lng,
                "altitude": new_alt,
                "relative_alt": new_alt,
                "ground_speed": flight_speed,
                "heading": heading,
                "yaw": heading,
                "airspeed": flight_speed * (0.8 + random.uniform(0, 0.4)),
                "battery_percentage": max(0, self.px4_status["battery_percentage"] - 0.02),
                "battery_current": 10.0 + random.uniform(-2, 2),
                "battery_voltage": 12.0 - (100 - self.px4_status["battery_percentage"]) * 0.01
            })
            
            # 更新电机速度和温度
            for i in range(4):
                self.px4_status["motors_speed"][i] = 1400 + random.uniform(-100, 100) + (flight_speed - 5) * 10
                self.px4_status["esc_temperature"][i] = min(75.0, 50.0 + random.uniform(-5, 5))
            
            # 更新油门
            self.px4_status["throttle"] = 0.5 + random.uniform(-0.1, 0.1)
            
            # 模拟姿态变化（与速度成正比）
            self.px4_status["roll"] = random.uniform(-5, 5) * (flight_speed / 8)
            self.px4_status["pitch"] = random.uniform(-2, 2) + (flight_speed / 8) * 2
            
            # 更新CPU温度和负载
            self.px4_status["cpu_temperature"] = 45.0 + random.uniform(-5, 5)
            self.px4_status["cpu_usage"] = 30.0 + random.uniform(-5, 10)
            
            # 模拟振动
            self.px4_status["vibration_x"] = 0.1 + random.uniform(0, 0.2) * (flight_speed / 8)
            self.px4_status["vibration_y"] = 0.1 + random.uniform(0, 0.2) * (flight_speed / 8)
            self.px4_status["vibration_z"] = 0.1 + random.uniform(0, 0.3) * (flight_speed / 8)
            
            # 向后台发送数据
            if self.current_task:
                self.send_flight_data(self.current_task['id'])
            
            time.sleep(self.update_interval)
        
        logger.info(f"[PX4] 到达航点 {self.current_waypoint_index + 1}")
        return True
    
    def start_mission(self, task):
        """开始执行任务"""
        logger.info(f"[PX4] 开始执行任务: {task['name']} (ID: {task['id']})")
        
        # 设置当前任务
        self.current_task = task
        self.is_flying = True
        
        # 获取任务航线
        routes = self.get_task_routes(task['id'])
        
        if not routes:
            logger.error("[PX4] 任务没有关联的航线")
            self.update_task_status(task['id'], "failed")
            self.is_flying = False
            return False
        
        # 设置当前航线（取第一个航线）
        self.current_route = routes[0]
        waypoints = self.current_route.get("waypoints", [])
        
        if not waypoints:
            logger.error("[PX4] 航线没有航点信息")
            self.update_task_status(task['id'], "failed")
            self.is_flying = False
            return False
        
        # 更新任务状态为执行中
        if not self.update_task_status(task['id'], "in_progress", 0):
            self.is_flying = False
            return False
        
        # 初始化飞行状态
        self.current_waypoint_index = 0
        self.px4_status["mission_item_total"] = len(waypoints)
        self.px4_status["mission_item_current"] = 0
        
        # 初始化位置到第一个航点附近
        first_waypoint = waypoints[0]
        self.px4_status.update({
            "latitude": first_waypoint["lat"] + random.uniform(-0.0001, 0.0001),
            "longitude": first_waypoint["lng"] + random.uniform(-0.0001, 0.0001),
            "altitude": 0.0,
            "relative_alt": 0.0,
            "ground_speed": 0.0,
            "heading": 0.0,
            "battery_percentage": 100
        })
        
        # 更新无人机状态
        self.update_drone_status()
        
        try:
            # 执行起飞
            if not self.takeoff(waypoints[0]["alt"]):
                raise Exception("起飞失败")
            
            # 执行每个航点
            total_waypoints = len(waypoints)
            for i in range(total_waypoints):
                self.current_waypoint_index = i
                
                # 执行当前航点
                if not self.execute_mission_item(waypoints[i]):
                    raise Exception(f"执行航点 {i+1} 失败")
                
                # 更新任务进度
                progress = int(((i + 1) / total_waypoints) * 100)
                self.update_task_status(task['id'], "in_progress", progress)
                
                # 如果不是最后一个航点，短暂悬停
                if i < total_waypoints - 1:
                    logger.info(f"[PX4] 悬停等待下一个航点...")
                    self.set_flight_mode("AUTO.LOITER")
                    self.px4_status["status_text"] = "Loitering"
                    
                    # 悬停2秒
                    hover_time = 2.0
                    start_time = time.time()
                    while time.time() - start_time < hover_time:
                        # 更新电池和状态
                        self.px4_status["battery_percentage"] = max(0, self.px4_status["battery_percentage"] - 0.01)
                        self.px4_status["roll"] = random.uniform(-1, 1)
                        self.px4_status["pitch"] = random.uniform(-1, 1)
                        
                        # 发送飞行数据
                        self.send_flight_data(task['id'])
                        time.sleep(self.update_interval)
            
            # 执行降落
            if not self.land():
                raise Exception("降落失败")
            
            # 任务完成
            self.update_task_status(task['id'], "completed", 100)
            logger.info(f"[PX4] 任务执行成功: {task['name']}")
            return True
            
        except Exception as e:
            logger.error(f"[PX4] 任务执行异常: {str(e)}")
            # 尝试降落
            try:
                self.land()
            except:
                pass
            
            # 更新任务状态为失败
            self.update_task_status(task['id'], "failed")
            return False
        finally:
            # 重置状态
            self.is_flying = False
            self.update_drone_status()
    
    async def run(self):
        """运行PX4模拟器"""
        logger.info("[PX4] PX4飞控模拟器启动")
        
        # 初始化状态
        self.px4_status["status_text"] = "PX4 initialized and ready"
        
        last_status_update = time.time()
        
        while True:
            try:
                # 定期更新无人机状态
                current_time = time.time()
                if current_time - last_status_update > 1.0:  # 每秒更新一次
                    self.update_drone_status()
                    last_status_update = current_time
                
                # 如果没有正在执行的任务，尝试获取新任务
                if not self.is_flying:
                    task = self.get_pending_tasks()
                    if task:
                        logger.info(f"[PX4] 发现新任务: {task['name']}")
                        # 阻塞执行任务
                        self.start_mission(task)
                    else:
                        # 随机模拟一些状态变化
                        self.px4_status["battery_percentage"] = min(100, self.px4_status["battery_percentage"] + random.uniform(0, 0.05))
                        self.px4_status["cpu_temperature"] = 40.0 + random.uniform(-2, 2)
                        
                        await asyncio.sleep(self.task_check_interval)
                
                # 小延迟避免CPU占用过高
                await asyncio.sleep(0.1)
                
            except Exception as e:
                logger.error(f"[PX4] 模拟器运行异常: {str(e)}")
                if self.is_flying:
                    try:
                        self.land()
                    except:
                        pass
                    self.is_flying = False
                await asyncio.sleep(5)

# 主程序
if __name__ == "__main__":
    simulator = PX4ControlSimulator()
    try:
        logger.info("PX4飞控控制模拟程序启动")
        logger.info("等待任务...")
        asyncio.run(simulator.run())
    except KeyboardInterrupt:
        logger.info("PX4模拟器已停止")
        # 尝试安全降落和上锁
        if simulator.is_flying:
            try:
                simulator.land()
            except:
                pass